package code.oldCode.feishuSpecializedTraining.dynamic;

import utils.TreeNode;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

public class MyDP6 {
    // 198. 打家劫舍
    public int rob(int[] nums) {
        // 1 <= nums.length <= 100
        int len = nums.length;
        if (len == 1)
            return nums[0];
        int[] dp = new int[len];
        dp[0] = nums[0];
        dp[1] = Math.max(nums[0], nums[1]);
        for (int i = 2; i < len; i++) {
            dp[i] = Math.max(dp[i - 1], dp[i - 2] + nums[i]);
        }
        return dp[len - 1];
    }

    // 213. 打家劫舍 II
    public int rob2(int[] nums) {
        // 1 <= nums.length <= 100
        int len = nums.length;
        if (len <= 3)
            return Arrays.stream(nums).max().getAsInt();
        // 由于问题发生在头尾，所以设置两个状态，一个是选择头，一个不选择头
        int[] dpWith0 = new int[len];
        int[] dpWithout0 = new int[len];
        dpWith0[0] = nums[0];
        dpWithout0[0] = 0;
        dpWith0[1] = nums[0];
        dpWithout0[1] = nums[1];
        for (int i = 2; i < len; i++) {
            if (i == len - 1)
                dpWith0[i] = dpWith0[i - 1];
            else
                dpWith0[i] = Math.max(dpWith0[i - 1], dpWith0[i - 2] + nums[i]);
            dpWithout0[i] = Math.max(dpWithout0[i - 1], dpWithout0[i - 2] + nums[i]);
        }
        // 这两种情况选最大值
        return Math.max(dpWith0[len - 1], dpWithout0[len - 1]);
    }

    // 337. 打家劫舍 III
    public int rob3(TreeNode root) {
        // 树形dp
        // 我用了一个hashmap，如果用两个map表示选取该节点和不选取该节点会更简单些
        Map<TreeNode, Integer> map = new HashMap<>();
        map.put(null, 0);
        robPostorder(map, root);
        return map.get(root);
    }

    private void robPostorder(Map<TreeNode, Integer> map, TreeNode root) {
        if (root == null) {
            return;
        }
        robPostorder(map, root.left);
        robPostorder(map, root.right);
        int noChoose = 0, choose = 0;
        if (root.left != null && root.right != null) {
            noChoose = map.get(root.left) + map.get(root.right);
            choose = root.val + map.get(root.left.left) + map.get(root.left.right) + map.get(root.right.left) + map.get(root.right.right);
        } else if (root.left != null) {
            noChoose = map.get(root.left);
            choose = root.val + map.get(root.left.left) + map.get(root.left.right);
        } else if (root.right != null) {
            noChoose = map.get(root.right);
            choose = root.val + map.get(root.right.left) + map.get(root.right.right);
        } else {
            choose = root.val;
        }
        map.put(root, Math.max(noChoose, choose));
    }

    public int rob3_twoMap(TreeNode root) {
        // 双map写法，表示选这个和不选这个，其中map其实可以优化掉，递归函数改为传递数组
        Map<TreeNode, Integer> mapChoose = new HashMap<>(), mapNo = new HashMap<>();
        mapChoose.put(null, 0);
        mapNo.put(null, 0);
        robPostorder_twoMap(mapChoose, mapNo, root);
        return Math.max(mapChoose.get(root), mapNo.get(root));
    }

    private void robPostorder_twoMap(Map<TreeNode, Integer> mapChoose, Map<TreeNode, Integer> mapNo, TreeNode root) {
        if (root == null)
            return;
        robPostorder_twoMap(mapChoose, mapNo, root.left);
        robPostorder_twoMap(mapChoose, mapNo, root.right);
        mapChoose.put(root, root.val + mapNo.get(root.left) + mapNo.get(root.right));
        // 两个map表示的时候注意，不选的情况下，子节点可以选也可以不选
        mapNo.put(root, Math.max(mapChoose.get(root.left), mapNo.get(root.left)) + Math.max(mapChoose.get(root.right), mapNo.get(root.right)));
    }

    public static void main(String[] args) {
        MyDP6 m = new MyDP6();
        System.out.println(m.rob2(new int[]{1, 2, 3, 5}));
    }
}
